import mxnet as mx

from .resnet_v1_101_rfcn_dcn import resnet_v1_101_rfcn_dcn


class DCNResnet101(mx.gluon.nn.HybridBlock):
    def __init__(self, *args, **kwargs):
        super(DCNResnet101, self).__init__(*args, **kwargs)
        data = mx.sym.var(name="data")
        Sym = resnet_v1_101_rfcn_dcn()
        res4b22_relu = Sym.get_resnet_v1_conv4(data)
        res5c_relu = Sym.get_resnet_v1_conv5(res4b22_relu)
        sym = mx.sym.Group([res4b22_relu, res5c_relu])

        with self.name_scope():
            self.features = mx.gluon.nn.SymbolBlock(
                sym,
                [data]
            )

    def hybrid_forward(self, F, x, *args, **kwargs):
        return self.features(x)
